import socket
import struct
import numpy as np
from scipy.io import wavfile
import datetime
import os
import time
import wave

# ================= CONFIGURATION =================
HOST = '0.0.0.0'
PORT = 50005
CHUNK_SIZE = 512  # Size for playback chunks

# Audio parameters (Must match Pico side exactly)
SAMPLE_RATE = 16000

def save_wav(pcm_data, filename):
    """
    User provided logic: No noise reduction, just bit shifting.
    """
    print(f"Processing and saving the audio {filename} ...")

    raw_data = np.frombuffer(pcm_data, dtype='<i4')
    
    if len(raw_data) == 0: return

    shift_amount = 16 
    
    # 32-bit to 16-bit conversion
    scaled_data = raw_data >> shift_amount
    scaled_data = np.clip(scaled_data, -32768, 32767)
    pcm_16bit = scaled_data.astype(np.int16)
    
    wavfile.write(filename, SAMPLE_RATE, pcm_16bit)
    
    full_path = os.path.abspath(filename)
    print(f"Successfully saved to: {full_path}")

def stream_audio_back(target_ip, target_port, filename):
    """
    Reads the saved WAV file and streams it back to the Pico via UDP.
    Includes pacing to prevent buffer overflow on the MCU.
    """
    if not os.path.exists(filename): 
        print("Error: File not found for playback.")
        return

    print(f"[Playback] Streaming back to {target_ip}:{target_port} ...")
    
    wf = wave.open(filename, 'rb')
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    
    try:
        # Read first chunk
        data = wf.readframes(CHUNK_SIZE // 2)
        
        while data:
            # Send UDP packet
            sock.sendto(data, (target_ip, target_port))
            
            # Read next chunk
            data = wf.readframes(CHUNK_SIZE // 2)
            
            # === PACING CONTROL ===
            # 16kHz * 16bit = 32000 bytes/sec
            # 512 bytes / 32000 = 0.016s
            # Sleep slightly less to keep buffer full but not overflowing
            time.sleep(0.015) 

        # Send End of Stream signal
        sock.sendto(b'END_OF_STREAM', (target_ip, target_port))
        print("[Playback] Transmission complete.")
        
    except Exception as e:
        print(f"Stream Error: {e}")
    finally:
        sock.close()
        wf.close()

def start_server():
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((HOST, PORT))
    print(f"Echo Server started on {HOST}:{PORT}")
    print("Logic: Record -> Save (No Filter) -> Auto Replay")
    print("Please press the button on the Pico to start...")

    frames = []
    is_recording = False
    client_addr = None # To store Pico's IP automatically

    try:
        while True:
            # Receive data
            data, addr = sock.recvfrom(4096)

            # Protocol logic
            if b'START' in data:
                print(f"[{datetime.datetime.now().strftime('%H:%M:%S')}] Start signal from {addr}, recording...")
                frames = []
                is_recording = True
                client_addr = addr # Remember who sent the audio
                
            elif b'STOP' in data:
                print("Stop signal received.")
                is_recording = False
                
                if len(frames) > 0:
                    # 1. Combine data
                    full_data = b''.join(frames)
                    
                    # 2. Generate filename
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    filename = f"mic_test_{timestamp}.wav"
                    
                    # 3. Save (Using YOUR strict logic)
                    save_wav(full_data, filename)
                    
                    # 4. Auto Echo (Playback)
                    if client_addr:
                        stream_audio_back(client_addr[0], client_addr[1], filename)
                        
                    print("--- Ready for next turn ---")
                else:
                    print("Warning: No audio data received")

            elif is_recording:
                frames.append(data)
                if len(frames) % 100 == 0:
                    print(".", end="", flush=True)

    except KeyboardInterrupt:
        print("\nServer stopped")
    finally:
        sock.close()

if __name__ == '__main__':
    start_server()